import math
import torch

import numpy as np
import torch.nn.functional as F

from PIL import Image
from tqdm import tqdm
from .floating_region import FloatingRegionScore
from .spatial_purity import SpatialPurity
from .region_relate_entropy import RegionRelateEntropyScore


def pixel_selection(logit, active_mask, active_indicator, active_selected, ground_truth, num_classes, active_pixels, mask_radius, select_fn=None):
    if select_fn is None:
        calculate_purity = SpatialPurity(in_channels=num_classes, size=2 * mask_radius + 1).cuda()
    else:
        calculate_purity = select_fn
    # p = torch.softmax(logit, dim=0)
    #
    # entropy = torch.sum(-p * torch.log(p + 1e-6), dim=0)
    #
    # pseudo_label = torch.argmax(p, dim=0)
    # one_hot = F.one_hot(pseudo_label, num_classes=num_classes).float()
    # one_hot = one_hot.permute((2, 0, 1)).unsqueeze(dim=0)
    # purity = calculate_purity(one_hot).squeeze(dim=0).squeeze(dim=0)
    # score = entropy * purity
    score = calculate_purity(logit)
    score[active_indicator] = -float('inf')

    for pixel in range(active_pixels):
        values, indices_h = torch.max(score, dim=0)
        _, indices_w = torch.max(values, dim=0)
        w = indices_w.item()
        h = indices_h[w].item()

        start_w = w - mask_radius if w - mask_radius >= 0 else 0
        start_h = h - mask_radius if h - mask_radius >= 0 else 0
        end_w = w + mask_radius + 1
        end_h = h + mask_radius + 1

        # mask out
        score[start_h:end_h, start_w:end_w] = -float('inf')
        active_indicator[start_h:end_h, start_w:end_w] = True
        active_selected[h, w] = True

        # active sampling
        active_mask[h, w] = ground_truth[h, w]
    return active_mask, active_indicator, active_selected


def region_selection(logit, active_mask, active_indicator, active_selected, ground_truth, num_classes, active_ratio, mask_radius, select_fn=None):
    if select_fn is None:
        floating_region_score = FloatingRegionScore(in_channels=num_classes, size=2 * mask_radius + 1).cuda()
    else:
        floating_region_score = select_fn
    per_region_pixels = (2 * mask_radius + 1) ** 2
    active_radius = mask_radius
    mask_radius = mask_radius * 2
    # active_ratio = cfg.ACTIVE.RATIO # / cfg.ACTIVE.SELECT_EPOCH
    # h, w = cfg.INPUT.TARGET_INPUT_SIZE_TRAIN
    h, w = logit.shape[-2:]
    num_pixel_cur = h * w
    active_regions = math.ceil(num_pixel_cur * active_ratio / per_region_pixels)
    # print('active_regions:{},\tper_region_pixels:{},\tactive_ratio:{}'.format(active_regions, per_region_pixels, active_ratio))

    score, purity, entropy = floating_region_score(logit)

    score[active_indicator] = -float('inf')

    for pixel in range(active_regions):
        values, indices_h = torch.max(score, dim=0)
        _, indices_w = torch.max(values, dim=0)
        w = indices_w.item()
        h = indices_h[w].item()

        active_start_w = w - active_radius if w - active_radius >= 0 else 0
        active_start_h = h - active_radius if h - active_radius >= 0 else 0
        active_end_w = w + active_radius + 1
        active_end_h = h + active_radius + 1

        mask_start_w = w - mask_radius if w - mask_radius >= 0 else 0
        mask_start_h = h - mask_radius if h - mask_radius >= 0 else 0
        mask_end_w = w + mask_radius + 1
        mask_end_h = h + mask_radius + 1

        # mask out
        score[mask_start_h:mask_end_h, mask_start_w:mask_end_w] = -float('inf')
        active_indicator[mask_start_h:mask_end_h, mask_start_w:mask_end_w] = True
        active_selected[active_start_h:active_end_h, active_start_w:active_end_w] = True

        # active sampling
        active_mask[active_start_h:active_end_h, active_start_w:active_end_w] = \
            ground_truth[active_start_h:active_end_h, active_start_w:active_end_w]
    return active_mask, active_indicator, active_selected


def softmax_pixel_selection(logit, active_mask, active_indicator, active_selected, ground_truth, num_classes, active_pixels, mask_radius, select_fn=None):
    score = torch.max(torch.softmax(logit, dim=0))[0].unsqueeze(dim=0).unsqueeze(dim=0)
    score[active_indicator] = -float('inf')

    for pixel in range(active_pixels):
        values, indices_h = torch.max(score, dim=0)
        _, indices_w = torch.max(values, dim=0)
        w = indices_w.item()
        h = indices_h[w].item()

        start_w = w - mask_radius if w - mask_radius >= 0 else 0
        start_h = h - mask_radius if h - mask_radius >= 0 else 0
        end_w = w + mask_radius + 1
        end_h = h + mask_radius + 1

        # mask out
        score[start_h:end_h, start_w:end_w] = -float('inf')
        active_indicator[start_h:end_h, start_w:end_w] = True
        active_selected[h, w] = True

        # active sampling
        active_mask[h, w] = ground_truth[h, w]
    return active_mask, active_indicator, active_selected


def softmax_region_selection(logit, active_mask, active_indicator, active_selected, ground_truth, num_classes, active_ratio, mask_radius, select_fn=None):
    per_region_pixels = (2 * mask_radius + 1) ** 2
    active_radius = mask_radius
    mask_radius = mask_radius * 2
    # active_ratio = cfg.ACTIVE.RATIO # / cfg.ACTIVE.SELECT_EPOCH
    # h, w = cfg.INPUT.TARGET_INPUT_SIZE_TRAIN
    h, w = logit.shape[-2:]
    num_pixel_cur = h * w
    active_regions = math.ceil(num_pixel_cur * active_ratio / per_region_pixels)
    # print('active_regions:{},\tper_region_pixels:{},\tactive_ratio:{}'.format(active_regions, per_region_pixels, active_ratio))

    score = torch.max(torch.softmax(logit, dim=0))[0].unsqueeze(dim=0).unsqueeze(dim=0)

    score[active_indicator] = -float('inf')

    for pixel in range(active_regions):
        values, indices_h = torch.max(score, dim=0)
        _, indices_w = torch.max(values, dim=0)
        w = indices_w.item()
        h = indices_h[w].item()

        active_start_w = w - active_radius if w - active_radius >= 0 else 0
        active_start_h = h - active_radius if h - active_radius >= 0 else 0
        active_end_w = w + active_radius + 1
        active_end_h = h + active_radius + 1

        mask_start_w = w - mask_radius if w - mask_radius >= 0 else 0
        mask_start_h = h - mask_radius if h - mask_radius >= 0 else 0
        mask_end_w = w + mask_radius + 1
        mask_end_h = h + mask_radius + 1

        # mask out
        score[mask_start_h:mask_end_h, mask_start_w:mask_end_w] = -float('inf')
        active_indicator[mask_start_h:mask_end_h, mask_start_w:mask_end_w] = True
        active_selected[active_start_h:active_end_h, active_start_w:active_end_w] = True

        # active sampling
        active_mask[active_start_h:active_end_h, active_start_w:active_end_w] = \
            ground_truth[active_start_h:active_end_h, active_start_w:active_end_w]
    return active_mask, active_indicator, active_selected


def ent_pixel_selection(logit, active_mask, active_indicator, active_selected, ground_truth, num_classes, active_pixels, mask_radius, select_fn=None):
    p = torch.softmax(logit, dim=0)
    score = torch.sum(-p * torch.log(p + 1e-6), dim=0).unsqueeze(dim=0).unsqueeze(dim=0)
    score[active_indicator] = -float('inf')

    for pixel in range(active_pixels):
        values, indices_h = torch.max(score, dim=0)
        _, indices_w = torch.max(values, dim=0)
        w = indices_w.item()
        h = indices_h[w].item()

        start_w = w - mask_radius if w - mask_radius >= 0 else 0
        start_h = h - mask_radius if h - mask_radius >= 0 else 0
        end_w = w + mask_radius + 1
        end_h = h + mask_radius + 1

        # mask out
        score[start_h:end_h, start_w:end_w] = -float('inf')
        active_indicator[start_h:end_h, start_w:end_w] = True
        active_selected[h, w] = True

        # active sampling
        active_mask[h, w] = ground_truth[h, w]
    return active_mask, active_indicator, active_selected


def ent_region_selection(logit, active_mask, active_indicator, active_selected, ground_truth, num_classes, active_ratio, mask_radius, select_fn=None):
    p = torch.softmax(logit, dim=0)
    score = torch.sum(-p * torch.log(p + 1e-6), dim=0).unsqueeze(dim=0).unsqueeze(dim=0)

    per_region_pixels = (2 * mask_radius + 1) ** 2
    active_radius = mask_radius
    mask_radius = mask_radius * 2
    # active_ratio = cfg.ACTIVE.RATIO # / cfg.ACTIVE.SELECT_EPOCH
    # h, w = cfg.INPUT.TARGET_INPUT_SIZE_TRAIN
    h, w = logit.shape[-2:]
    num_pixel_cur = h * w
    active_regions = math.ceil(num_pixel_cur * active_ratio / per_region_pixels)
    # print('active_regions:{},\tper_region_pixels:{},\tactive_ratio:{}'.format(active_regions, per_region_pixels, active_ratio))

    # score, purity, entropy = floating_region_score(logit)

    score[active_indicator] = -float('inf')

    for pixel in range(active_regions):
        values, indices_h = torch.max(score, dim=0)
        _, indices_w = torch.max(values, dim=0)
        w = indices_w.item()
        h = indices_h[w].item()

        active_start_w = w - active_radius if w - active_radius >= 0 else 0
        active_start_h = h - active_radius if h - active_radius >= 0 else 0
        active_end_w = w + active_radius + 1
        active_end_h = h + active_radius + 1

        mask_start_w = w - mask_radius if w - mask_radius >= 0 else 0
        mask_start_h = h - mask_radius if h - mask_radius >= 0 else 0
        mask_end_w = w + mask_radius + 1
        mask_end_h = h + mask_radius + 1

        # mask out
        score[mask_start_h:mask_end_h, mask_start_w:mask_end_w] = -float('inf')
        active_indicator[mask_start_h:mask_end_h, mask_start_w:mask_end_w] = True
        active_selected[active_start_h:active_end_h, active_start_w:active_end_w] = True

        # active sampling
        active_mask[active_start_h:active_end_h, active_start_w:active_end_w] = \
            ground_truth[active_start_h:active_end_h, active_start_w:active_end_w]
    return active_mask, active_indicator, active_selected


def random_pixel_selection(logit, active_mask, active_indicator, active_selected, ground_truth, num_classes, active_pixels, mask_radius, select_fn=None):
    score = torch.randn(logit.shape[-2:])
    score[active_indicator] = -float('inf')

    for pixel in range(active_pixels):
        values, indices_h = torch.max(score, dim=0)
        _, indices_w = torch.max(values, dim=0)
        w = indices_w.item()
        h = indices_h[w].item()

        start_w = w - mask_radius if w - mask_radius >= 0 else 0
        start_h = h - mask_radius if h - mask_radius >= 0 else 0
        end_w = w + mask_radius + 1
        end_h = h + mask_radius + 1

        # mask out
        score[start_h:end_h, start_w:end_w] = -float('inf')
        active_indicator[start_h:end_h, start_w:end_w] = True
        active_selected[h, w] = True

        # active sampling
        active_mask[h, w] = ground_truth[h, w]
    return active_mask, active_indicator, active_selected


def random_region_selection(logit, active_mask, active_indicator, active_selected, ground_truth, num_classes, active_ratio, mask_radius, select_fn=None):
    per_region_pixels = (2 * mask_radius + 1) ** 2
    active_radius = mask_radius
    mask_radius = mask_radius * 2
    h, w = logit.shape[-2:]
    num_pixel_cur = h * w
    active_regions = math.ceil(num_pixel_cur * active_ratio / per_region_pixels)
    # score, purity, entropy = floating_region_score(logit)
    score = torch.randn(logit.shape[-2:])
    score[active_indicator] = -float('inf')

    for pixel in range(active_regions):
        values, indices_h = torch.max(score, dim=0)
        _, indices_w = torch.max(values, dim=0)
        w = indices_w.item()
        h = indices_h[w].item()

        active_start_w = w - active_radius if w - active_radius >= 0 else 0
        active_start_h = h - active_radius if h - active_radius >= 0 else 0
        active_end_w = w + active_radius + 1
        active_end_h = h + active_radius + 1

        mask_start_w = w - mask_radius if w - mask_radius >= 0 else 0
        mask_start_h = h - mask_radius if h - mask_radius >= 0 else 0
        mask_end_w = w + mask_radius + 1
        mask_end_h = h + mask_radius + 1

        # mask out
        score[mask_start_h:mask_end_h, mask_start_w:mask_end_w] = -float('inf')
        active_indicator[mask_start_h:mask_end_h, mask_start_w:mask_end_w] = True
        active_selected[active_start_h:active_end_h, active_start_w:active_end_w] = True

        # active sampling
        active_mask[active_start_h:active_end_h, active_start_w:active_end_w] = \
            ground_truth[active_start_h:active_end_h, active_start_w:active_end_w]
    return active_mask, active_indicator, active_selected


def test_selection(logit, active_mask, active_indicator, active_selected, ground_truth, num_classes, active_ratio, mask_radius, select_fn=None):
    if select_fn is None:
        floating_region_score = RegionRelateEntropyScore(in_channels=num_classes, size=2 * mask_radius + 1).cuda()
    else:
        floating_region_score = select_fn
    per_region_pixels = (2 * mask_radius + 1) ** 2
    active_radius = mask_radius
    mask_radius = mask_radius * 2
    # active_ratio = cfg.ACTIVE.RATIO # / cfg.ACTIVE.SELECT_EPOCH
    # h, w = cfg.INPUT.TARGET_INPUT_SIZE_TRAIN
    h, w = logit.shape[-2:]
    num_pixel_cur = h * w
    active_regions = math.ceil(num_pixel_cur * active_ratio / per_region_pixels)
    # print('active_regions:{},\tper_region_pixels:{},\tactive_ratio:{}'.format(active_regions, per_region_pixels, active_ratio))

    score, purity, entropy = floating_region_score(logit)

    score[active_indicator] = -float('inf')
    score[ground_truth == 255] = -float('inf')

    for pixel in range(active_regions):
        values, indices_h = torch.max(score, dim=0)
        _, indices_w = torch.max(values, dim=0)
        w = indices_w.item()
        h = indices_h[w].item()

        active_start_w = w - active_radius if w - active_radius >= 0 else 0
        active_start_h = h - active_radius if h - active_radius >= 0 else 0
        active_end_w = w + active_radius + 1
        active_end_h = h + active_radius + 1

        mask_start_w = w - mask_radius if w - mask_radius >= 0 else 0
        mask_start_h = h - mask_radius if h - mask_radius >= 0 else 0
        mask_end_w = w + mask_radius + 1
        mask_end_h = h + mask_radius + 1

        # mask out
        score[mask_start_h:mask_end_h, mask_start_w:mask_end_w] = -float('inf')
        active_indicator[mask_start_h:mask_end_h, mask_start_w:mask_end_w] = True
        active_selected[active_start_h:active_end_h, active_start_w:active_end_w] = True

        # active sampling
        active_mask[active_start_h:active_end_h, active_start_w:active_end_w] = \
            ground_truth[active_start_h:active_end_h, active_start_w:active_end_w]
    return active_mask, active_indicator, active_selected
